import torch
from torch import nn

"""
Transformer
===========================__init___参数：=================================
d_model (int) – 编码器/解码器输入中预期的特征数量（默认值=512）。
nhead (int) – 多头注意力模型中的头数量（默认值=8）。
num_encoder_layers (int) – 编码器中的子编码器层数量（默认值=6）。
num_decoder_layers (int) – 解码器中的子解码器层数量（默认值=6）。
dim_feedforward (int) – 前馈网络模型的维度（默认值=2048）。
dropout (float) – dropout 值（默认值=0.1）。
activation (Union[str, Callable[[Tensor], Tensor]]) – 编码器/解码器中间层的激活函数，可以是字符串（“relu”或“gelu”）或一元可调用对象。默认值：relu
custom_encoder (Optional[Any]) – 自定义编码器（默认值=None）。
custom_decoder (Optional[Any]) – 自定义解码器（默认值=None）。
layer_norm_eps (float) – 层归一化组件中的 eps 值（默认值=1e-5）。
batch_first (bool) – 如果为 True，则输入和输出张量的形状为 (batch, seq, feature)。默认值：False (seq, batch, feature)。
norm_first (bool) – 如果为 True，编码器和解码器层将在其他注意力层和前馈操作之前执行 LayerNorm，否则在其后执行。默认值：False (之后)。
bias (bool) – 如果设置为 False，Linear 和 LayerNorm 层将不会学习加性偏置。默认值：True。
===========================forward参数：=================================
src (Tensor) – 编码器输入的序列（必需）。
tgt (Tensor) – 解码器输入的序列（必需）。
src_mask (Optional[Tensor]) – src 序列的加性掩码（可选）。
tgt_mask (Optional[Tensor]) – tgt 序列的加性掩码（可选）。
memory_mask (Optional[Tensor]) – 编码器输出的加性掩码（可选）。
src_key_padding_mask (Optional[Tensor]) – src 键的张量掩码，按批次划分（可选）。
tgt_key_padding_mask (Optional[Tensor]) – tgt 键的张量掩码，按批次划分（可选）。
memory_key_padding_mask (Optional[Tensor]) – memory 键的张量掩码，按批次划分（可选）。
src_is_causal (Optional[bool]) – 如果指定，则将因果掩码应用为 src_mask。默认值：None；尝试检测因果掩码。警告：src_is_causal 提供了一个提示，表明 src_mask 是因果掩码。提供不正确的提示可能导致执行错误，包括前向和后向兼容性问题。
tgt_is_causal (Optional[bool]) – 如果指定，则将因果掩码应用为 tgt_mask。默认值：None；尝试检测因果掩码。警告：tgt_is_causal 提供了一个提示，表明 tgt_mask 是因果掩码。提供不正确的提示可能导致执行错误，包括前向和后向兼容性问题。
memory_is_causal (bool) – 如果指定，则将因果掩码应用为 memory_mask。默认值：False。警告：memory_is_causal 提供了一个提示，表明 memory_mask 是因果掩码。提供不正确的提示可能导致执行错误，包括前向和后向兼容性问题。
"""
transformer_model = nn.Transformer(nhead=16, num_encoder_layers=12, batch_first=True)
src = torch.rand((32, 10, 512))
tgt = torch.rand((32, 20, 512))
out = transformer_model(src, tgt)
print(out.shape)
